'''
配置文件
'''

import os
import torch
import datetime
import logging
import sys
from easydict import EasyDict as edict

__C = edict()
args = __C

# Base
__C.gpu_id = 3
__C.num_workers = 4
__C.seed = 666
__C.do_train = True
__C.do_test = False
__C.batch_size = 32
__C.lr = 0.0001
__C.now = datetime.datetime.now()
__C.device = 'cuda' if torch.cuda.is_available() else 'cpu'
# __C.extract_algorithm = "StanfordCoreNLP"       # "StanfordCoreNLP" / "Spacy"/ "use_Spacy"
__C.extract_algorithm = "Spacy"
torch.cuda.set_device(__C.gpu_id)

# dataset options
__C.dataset = edict()
__C.dataset.name = 'twitter2017'
__C.dataset.type = 'none'
__C.dataset.data_dir = os.path.join('dataset',  __C.dataset.name)
__C.dataset.save_dir = os.path.join('result', __C.dataset.name, 'predict')
__C.dataset.ckpt_dir = os.path.join('result', __C.dataset.name, 'ckpt')
__C.dataset.log_dir = os.path.join('result', __C.dataset.name, 'log')

__C.dataset.predict_file = __C.dataset.save_dir + "/{}/epoch{}.txt"
__C.dataset.image_caption = os.path.join('captions', '{}_image_captions.json')
__C.dataset.text = os.path.join('texts', '{}_text.json')
__C.dataset.phrase = os.path.join('texts', '{}_{}_phrase.json')
__C.dataset.image = os.path.join('images', 'image_features_{}.pkl')
__C.dataset.image_77 = os.path.join('images', 'image_features_7_7_{}.pkl')
__C.dataset = dict(__C.dataset)


# 创建文件夹
if not os.path.exists(__C.dataset.save_dir):
    os.makedirs(__C.dataset.save_dir)
if not os.path.exists(__C.dataset.ckpt_dir):
    os.makedirs(__C.dataset.ckpt_dir)
if not os.path.exists(__C.dataset.log_dir):
    os.makedirs(__C.dataset.log_dir)

# logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)-8s %(message)s')
logFormatter = logging.Formatter('%(asctime)s %(levelname)-8s %(message)s')
rootLogger = logging.getLogger()
fileHandler = logging.FileHandler(
        os.path.join(__C.dataset.log_dir, 'stdout_{}.{}_{}.{}.log'.format(__C.now.month, __C.now.day, __C.now.hour, __C.now.minute)), 'w+')
fileHandler.setFormatter(logFormatter)
rootLogger.addHandler(fileHandler)
__C.logging = logging

# training options
__C.train = edict()
__C.train.max_epochs = 100
__C.train.max_word_length = 128
__C.train.bert_dim = 768
__C.train.module_dim = 1024
__C.train.interval = 20
__C.train.tag2idx = {
    "PAD": 0,
    "B-PER": 1,
    "I-PER": 2,
    "B-LOC": 3,
    "I-LOC": 4,
    "B-ORG": 5,
    "I-ORG": 6,
    "B-MISC": 7,
    "I-MISC": 8,
    "O": 9,
    "X": 10,
    "CLS": 11,
    "SEP": 12
}
__C.train.idx2tag = {str(idx): tag for tag, idx in __C.train.tag2idx.items()}
__C.train.pic_classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
           'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
           'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
           'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
           'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
           'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
           'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
           'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
           'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
           'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
           'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
           'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
           'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
           'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']
__C.train = dict(__C.train)
